package com.golang.gateway.filter;


import com.golang.gateway.utils.LogHelper;
import io.netty.buffer.UnpooledByteBufAllocator;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.reactivestreams.Publisher;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerCodecConfigurer;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.nio.charset.Charset;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicReference;


/**
 * 日志过滤器
 */

@Slf4j
@Component
public class GatewayLogFilter implements GlobalFilter, Ordered {

    @Autowired
    ServerCodecConfigurer codecConfig;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        try {
            ServerHttpRequest request = exchange.getRequest();
            ServerRequest serverRequest = ServerRequest.create(exchange, codecConfig.getReaders());
            URI requestUri = request.getURI();
            String uriQuery = requestUri.getQuery();
            HttpHeaders headers = request.getHeaders();
            MediaType mediaType = headers.getContentType();
            String schema = requestUri.getScheme();
            String method = request.getMethodValue().toUpperCase();

            if ((!"http".equals(schema) && !"https".equals(schema))) {
                return chain.filter(exchange);
            }
            // 原始请求体
            final AtomicReference<String> requestBody = new AtomicReference<>();
            if (Objects.nonNull(mediaType) && !mediaType.toString().startsWith(MediaType.APPLICATION_JSON.toString())) {
                requestBody.set("上传文件");
                return chain.filter(exchange);
            } else {
                if (method.equals("GET")) {
                    if (StringUtils.isNotBlank(uriQuery)) {
                        requestBody.set(uriQuery);
                    }
                } else if (headers.getContentLength() > 0) {
                    return serverRequest.bodyToMono(String.class).flatMap(body -> {
                        requestBody.set(body);
                        ServerHttpRequestDecorator requestDecorator = new ServerHttpRequestDecorator(exchange.getRequest()) {
                            @Override
                            public HttpHeaders getHeaders() {
                                HttpHeaders httpHeaders = new HttpHeaders();
                                httpHeaders.putAll(super.getHeaders());
                                return httpHeaders;
                            }

                            @Override
                            public Flux<DataBuffer> getBody() {
                                NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(new UnpooledByteBufAllocator(false));
                                DataBuffer bodyDataBuffer = nettyDataBufferFactory.wrap(body.getBytes());
                                return Flux.just(bodyDataBuffer);
                            }
                        };

                        ServerHttpResponseDecorator responseDecorator = getServerHttpResponseDecorator(exchange, requestBody);
                        return chain.filter(exchange.mutate()
                                .request(requestDecorator)
                                .response(responseDecorator)
                                .build());
                    });
                }
            }
            ServerHttpResponseDecorator decoratedResponse = getServerHttpResponseDecorator(exchange,
                    requestBody);
            return chain.filter(exchange.mutate()
                    .response(decoratedResponse)
                    .build());
        } catch (Exception e) {
            log.error("请求响应日志打印出现异常", e);
            return chain.filter(exchange);
        }
    }

    private ServerHttpResponseDecorator getServerHttpResponseDecorator(ServerWebExchange exchange, AtomicReference<String> requestBody) {

        ServerHttpResponse originalResponse = exchange.getResponse();
        DataBufferFactory bufferFactory = originalResponse.bufferFactory();
        HttpStatus httpStatus = originalResponse.getStatusCode();
        ServerHttpRequest request = exchange.getRequest();
        URI requestUri = request.getURI();
        String uriQuery = requestUri.getQuery();
        String url = requestUri.getPath() + (StringUtils.isNotBlank(uriQuery) ? "?" + uriQuery : "");
        String method = request.getMethodValue().toUpperCase();

        return new ServerHttpResponseDecorator(originalResponse) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                if (body instanceof Flux) {
                    Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                    return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                        DataBuffer join = bufferFactory.join(dataBuffers);
                        byte[] content = new byte[join.readableByteCount()];
                        join.read(content);
                        DataBufferUtils.release(join);
                        Charset charset = LogHelper.getMediaTypeCharset(originalResponse.getHeaders().getContentType());
                        String responseBody = new String(content, charset);
                        log.info("url:{}||method:{}||status:{}||request:{}||response:{}", url, method, httpStatus.value(), LogHelper.replaceAllBlank(requestBody.get()), responseBody);
                        return bufferFactory.wrap(content);
                    }));
                }
                return super.writeWith(body);
            }
        };
    }

    @Override
    public int getOrder() {
        return -2;
    }

}
